# AI Booster PowerShell 部署脚本
# 用于从Windows部署到Linux服务器

param(
    [Parameter(Mandatory=$false)]
    [string]$ServerIP,
    
    [Parameter(Mandatory=$false)]
    [string]$Username = "root",
    
    [Parameter(Mandatory=$false)]
    [int]$Port = 22,
    
    [Parameter(Mandatory=$false)]
    [string]$Command
)

# 颜色输出函数
function Write-ColorOutput {
    param(
        [string]$Message,
        [string]$Type = "Info"
    )
    
    switch ($Type) {
        "Success" { Write-Host "[SUCCESS] $Message" -ForegroundColor Green }
        "Error"   { Write-Host "[ERROR] $Message" -ForegroundColor Red }
        "Warning" { Write-Host "[WARNING] $Message" -ForegroundColor Yellow }
        "Info"    { Write-Host "[INFO] $Message" -ForegroundColor Blue }
        default   { Write-Host $Message }
    }
}

# 显示帮助信息
function Show-Help {
    Write-Host "AI Booster PowerShell 部署脚本" -ForegroundColor Cyan
    Write-Host ""
    Write-Host "使用方法:"
    Write-Host "  .\deploy.ps1 -ServerIP <IP> [-Username <用户名>] [-Port <端口>]"
    Write-Host ""
    Write-Host "参数说明:"
    Write-Host "  -ServerIP   目标服务器IP地址"
    Write-Host "  -Username   SSH用户名 (默认: root)"
    Write-Host "  -Port       SSH端口 (默认: 22)"
    Write-Host ""
    Write-Host "示例:"
    Write-Host "  .\deploy.ps1 -ServerIP 192.168.1.100"
    Write-Host "  .\deploy.ps1 -ServerIP 192.168.1.100 -Username ubuntu -Port 2222"
    Write-Host ""
    Write-Host "其他命令:"
    Write-Host "  .\deploy.ps1 -Command help"
    Write-Host "  .\deploy.ps1 -Command check"
    Write-Host ""
    Write-Host "注意事项:"
    Write-Host "  - 需要安装并配置SSH客户端（如OpenSSH或PuTTY）"
    Write-Host "  - 建议配置SSH密钥认证"
    Write-Host "  - 确保目标服务器可以通过SSH访问"
}

# 检查必需文件
function Test-RequiredFiles {
    $requiredFiles = @("install_ai_booster.sh", "manage_ai_booster.sh")
    $missingFiles = @()
    
    foreach ($file in $requiredFiles) {
        if (-not (Test-Path $file)) {
            $missingFiles += $file
        }
    }
    
    if ($missingFiles.Count -gt 0) {
        Write-ColorOutput "缺少必需文件:" "Error"
        foreach ($file in $missingFiles) {
            Write-Host "  - $file" -ForegroundColor Red
        }
        Write-ColorOutput "请确保在包含安装脚本的目录中运行此脚本" "Info"
        return $false
    }
    
    Write-ColorOutput "所有必需文件已找到" "Success"
    return $true
}

# 检查SSH客户端
function Test-SSHClient {
    $sshPath = Get-Command ssh -ErrorAction SilentlyContinue
    $scpPath = Get-Command scp -ErrorAction SilentlyContinue
    
    if (-not $sshPath) {
        Write-ColorOutput "未找到SSH客户端" "Error"
        Write-Host "请安装OpenSSH客户端或确保SSH在PATH中" -ForegroundColor Yellow
        Write-Host "Windows 10/11用户可以通过以下方式安装OpenSSH:"
        Write-Host "  设置 -> 应用 -> 可选功能 -> 添加功能 -> OpenSSH客户端"
        return $false
    }
    
    if (-not $scpPath) {
        Write-ColorOutput "未找到SCP客户端" "Error"
        return $false
    }
    
    Write-ColorOutput "SSH/SCP客户端检查通过" "Success"
    return $true
}

# 测试SSH连接
function Test-SSHConnection {
    param(
        [string]$ServerIP,
        [string]$Username,
        [int]$Port
    )
    
    Write-ColorOutput "正在测试SSH连接到 $Username@$ServerIP..." "Info"
    
    $testCommand = "ssh -p $Port -o ConnectTimeout=10 -o BatchMode=yes $Username@$ServerIP `"echo 'SSH连接测试成功'`""
    
    try {
        $result = Invoke-Expression $testCommand 2>$null
        if ($LASTEXITCODE -eq 0) {
            Write-ColorOutput "SSH密钥认证成功" "Success"
            return $true
        } else {
            Write-ColorOutput "SSH密钥认证失败，将使用密码认证" "Warning"
            return $false
        }
    } catch {
        Write-ColorOutput "SSH连接测试失败: $($_.Exception.Message)" "Error"
        return $false
    }
}

# 部署到服务器
function Deploy-ToServer {
    param(
        [string]$ServerIP,
        [string]$Username,
        [int]$Port
    )
    
    Write-ColorOutput "开始部署到服务器: $Username@$ServerIP`:$Port" "Info"
    
    # 检查SSH连接
    $sshKeyAuth = Test-SSHConnection -ServerIP $ServerIP -Username $Username -Port $Port
    if (-not $sshKeyAuth) {
        Write-ColorOutput "请准备输入SSH密码" "Info"
    }
    
    # 上传安装脚本
    Write-ColorOutput "正在上传 install_ai_booster.sh..." "Info"
    $scpCommand = "scp -P $Port install_ai_booster.sh $Username@$ServerIP`:/tmp/"
    
    try {
        Invoke-Expression $scpCommand
        if ($LASTEXITCODE -eq 0) {
            Write-ColorOutput "install_ai_booster.sh 上传成功" "Success"
        } else {
            Write-ColorOutput "install_ai_booster.sh 上传失败" "Error"
            return $false
        }
    } catch {
        Write-ColorOutput "上传失败: $($_.Exception.Message)" "Error"
        return $false
    }
    
    # 上传管理脚本
    Write-ColorOutput "正在上传 manage_ai_booster.sh..." "Info"
    $scpCommand = "scp -P $Port manage_ai_booster.sh $Username@$ServerIP`:/tmp/"
    
    try {
        Invoke-Expression $scpCommand
        if ($LASTEXITCODE -eq 0) {
            Write-ColorOutput "manage_ai_booster.sh 上传成功" "Success"
        } else {
            Write-ColorOutput "manage_ai_booster.sh 上传失败" "Error"
            return $false
        }
    } catch {
        Write-ColorOutput "上传失败: $($_.Exception.Message)" "Error"
        return $false
    }
    
    # 上传README（可选）
    if (Test-Path "README.md") {
        $scpCommand = "scp -P $Port README.md $Username@$ServerIP`:/tmp/"
        try {
            Invoke-Expression $scpCommand 2>$null
            Write-ColorOutput "README.md 上传成功" "Info"
        } catch {
            # 忽略README上传失败
        }
    }
    
    # 在远程服务器上设置权限
    Write-ColorOutput "正在远程服务器上设置权限..." "Info"
    
    $remoteCommands = @"
cd /tmp
chmod +x install_ai_booster.sh
chmod +x manage_ai_booster.sh
echo ""
echo "========================================="
echo "     文件上传完成！"
echo "========================================="
echo ""
echo "文件位置: /tmp/"
echo "  - install_ai_booster.sh  (安装脚本)"
echo "  - manage_ai_booster.sh   (管理脚本)"
echo "  - README.md             (说明文档)"
echo ""
echo "安装命令:"
echo "  sudo bash /tmp/install_ai_booster.sh"
echo ""
echo "管理命令:"
echo "  sudo bash /tmp/manage_ai_booster.sh config"
echo ""
echo "========================================="
"@
    
    $sshCommand = "ssh -p $Port $Username@$ServerIP `"$remoteCommands`""
    
    try {
        Invoke-Expression $sshCommand
        if ($LASTEXITCODE -eq 0) {
            Write-ColorOutput "远程部署完成！" "Success"
            Write-Host ""
            Write-ColorOutput "下一步操作:" "Info"
            Write-Host "  1. SSH登录到服务器: ssh -p $Port $Username@$ServerIP"
            Write-Host "  2. 运行安装命令: sudo bash /tmp/install_ai_booster.sh"
            Write-Host "  3. 开始使用: source /etc/ai_booster enable"
            return $true
        } else {
            Write-ColorOutput "远程命令执行失败" "Error"
            return $false
        }
    } catch {
        Write-ColorOutput "远程命令执行失败: $($_.Exception.Message)" "Error"
        return $false
    }
}

# 批量部署
function Deploy-Batch {
    param(
        [string]$ServersFile = "servers.txt"
    )
    
    if (-not (Test-Path $ServersFile)) {
        Write-ColorOutput "服务器列表文件不存在: $ServersFile" "Error"
        Write-ColorOutput "请创建服务器列表文件或使用 -Command create-template 创建模板" "Info"
        return
    }
    
    Write-ColorOutput "开始批量部署，读取服务器列表: $ServersFile" "Info"
    Write-ColorOutput "文件格式: ip:username:port (每行一个服务器)" "Warning"
    
    $servers = Get-Content $ServersFile | Where-Object { $_ -notmatch '^\s*#' -and $_ -notmatch '^\s*$' }
    $successCount = 0
    $failCount = 0
    
    foreach ($server in $servers) {
        $parts = $server -split ':'
        $ip = $parts[0].Trim()
        $user = if ($parts[1]) { $parts[1].Trim() } else { "root" }
        $portNum = if ($parts[2]) { [int]$parts[2].Trim() } else { 22 }
        
        Write-Host ""
        Write-Host "=========================================" -ForegroundColor Cyan
        Write-Host "部署到服务器: $ip" -ForegroundColor Cyan
        Write-Host "=========================================" -ForegroundColor Cyan
        
        if (Deploy-ToServer -ServerIP $ip -Username $user -Port $portNum) {
            $successCount++
            Write-ColorOutput "服务器 $ip 部署成功" "Success"
        } else {
            $failCount++
            Write-ColorOutput "服务器 $ip 部署失败" "Error"
        }
        
        Start-Sleep -Seconds 2  # 避免连接过快
    }
    
    Write-Host ""
    Write-Host "=========================================" -ForegroundColor Cyan
    Write-Host "批量部署完成！" -ForegroundColor Cyan
    Write-Host "=========================================" -ForegroundColor Cyan
    Write-Host "成功: $successCount 台服务器" -ForegroundColor Green
    Write-Host "失败: $failCount 台服务器" -ForegroundColor Red
    Write-Host "总计: $($successCount + $failCount) 台服务器"
}

# 创建服务器列表模板
function New-ServerListTemplate {
    $template = @"
# AI Booster 服务器列表
# 格式: IP:用户名:端口
# 示例:
# 192.168.1.100:root:22
# 192.168.1.101:ubuntu:2222
# 192.168.1.102::   (使用默认用户名root和端口22)

"@
    
    $template | Out-File -FilePath "servers.txt" -Encoding UTF8
    Write-ColorOutput "服务器列表模板已创建: servers.txt" "Success"
    Write-ColorOutput "请编辑 servers.txt 文件，添加您的服务器信息" "Info"
}

# 主函数
function Main {
    # 处理命令
    switch ($Command) {
        "help" {
            Show-Help
            return
        }
        "check" {
            Write-ColorOutput "正在检查环境..." "Info"
            $filesOk = Test-RequiredFiles
            $sshOk = Test-SSHClient
            
            if ($filesOk -and $sshOk) {
                Write-ColorOutput "环境检查通过，可以开始部署" "Success"
            } else {
                Write-ColorOutput "环境检查失败，请解决上述问题后重试" "Error"
            }
            return
        }
        "batch" {
            if (-not (Test-RequiredFiles) -or -not (Test-SSHClient)) {
                return
            }
            Deploy-Batch
            return
        }
        "create-template" {
            New-ServerListTemplate
            return
        }
    }
    
    # 如果没有提供服务器IP且没有命令，显示帮助
    if (-not $ServerIP) {
        Show-Help
        return
    }
    
    # 验证IP地址格式
    if ($ServerIP -notmatch '^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$') {
        Write-ColorOutput "无效的IP地址: $ServerIP" "Error"
        return
    }
    
    # 检查环境
    if (-not (Test-RequiredFiles) -or -not (Test-SSHClient)) {
        return
    }
    
    # 部署到单个服务器
    Deploy-ToServer -ServerIP $ServerIP -Username $Username -Port $Port
}

# 运行主函数
Main
